#include <algorithm>
#include <functional>
#include <map>
#include <set>
#include <vector>
#include <cfloat>


#include "caffe/layer.hpp"
#include "caffe/layers/yolov2_detector_layer.hpp"
#include "caffe/net.hpp"

namespace caffe {

int entry_index(int w, int h, int c, int coords, int classes,
		int batch, int location, int entry) {
    int n =   location / (w*h);
    int loc = location % (w*h);
	int outputs = c * h * w;
    return batch*outputs + n*w*h*(coords+classes+1) + entry*w*h + loc;
}

template <typename Dtype>
void softmax(const Dtype *input, int n, Dtype temp, int stride, Dtype *output) {
    int i;
    float sum = 0;
    float largest = -FLT_MAX;
    for(i = 0; i < n; ++i){
        if(input[i*stride] > largest) largest = input[i*stride];
    }
    for(i = 0; i < n; ++i){
        float e = exp(input[i*stride]/temp - largest/temp);
        sum += e;
        output[i*stride] = e;
    }
    for(i = 0; i < n; ++i){
        output[i*stride] /= sum;
    }
}

template <typename Dtype>
void softmax_cpu(const Dtype *input, int n, int batch, int batch_offset, int groups,
		int group_offset, int stride, Dtype temp, Dtype *output) {
    int g, b;
    for(b = 0; b < batch; ++b){
        for(g = 0; g < groups; ++g){
            softmax(input + b*batch_offset + g*group_offset, n, temp, stride,
					output + b*batch_offset + g*group_offset);
        }
    }
}

template <typename Dtype>
static inline Dtype logistic_activate(Dtype x){return 1./(1. + exp(-x));}

template <typename Dtype>
static inline void logistic_activate_array(Dtype *x, const int n) {
	for (int i = 0; i < n; ++i) {
		x[i] = logistic_activate(x[i]);
	}
}

template <typename Dtype>
void YOLOv2DetectorLayer<Dtype>::LayerSetUp(const vector<Blob<Dtype>*>& bottom,
    const vector<Blob<Dtype>*>& top) {
  const YOLOv2DetectorParameter& param = this->layer_param_.yolov2_detector_param();
  CHECK_EQ(param.anchor_w_size(), param.anchor_h_size())
    << "anchor_w and anchor_h must be the sample dimension.";
  anchor_num_ = param.anchor_w_size();
  for (int i = 0; i < anchor_num_; ++i) {
    float w = param.anchor_w(i);
    float h = param.anchor_h(i);
    std::vector<float> anchor_size(2);
    anchor_size[0] = w;
    anchor_size[1] = h;
    anchors_.push_back(anchor_size);
  }
  class_num_ = param.classes();
  coords_ = param.coords();
}

template <typename Dtype>
void YOLOv2DetectorLayer<Dtype>::Reshape(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top) {
  // Initialize with the first blob.
  vector<int> top_shape = bottom[0]->shape();
  top[0]->Reshape(top_shape);
}

template <typename Dtype>
void YOLOv2DetectorLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top) {
  //const int size = coords_ + 1 + class_num_;
  for (int i = 0; i < bottom.size(); ++i) {
  	Dtype* top_data = top[i]->mutable_cpu_data();
    const Dtype* bottom_data = bottom[i]->cpu_data();
	memcpy(top_data, bottom_data, bottom[i]->count() * sizeof(Dtype));
	for (int b = 0; b < bottom[i]->num(); ++b) {
		for (int n = 0; n < anchor_num_; ++n) {
			// tx, ty
			int index = entry_index(bottom[i]->width(), bottom[i]->height(),
					bottom[i]->channels(), coords_, class_num_,
					b, n * bottom[i]->count(2), 0);
			logistic_activate_array(top_data + index,
					bottom[i]->height() * bottom[i]->width() * 2);
			// scale, logistic activate to
			index = entry_index(bottom[i]->width(), bottom[i]->height(),
					bottom[i]->channels(), coords_, class_num_,
					b, n * bottom[i]->count(2), 4);
			logistic_activate_array(top_data + index,
					bottom[i]->height() * bottom[i]->width());
		}
	}
	// softmax label
	int index = entry_index(bottom[i]->width(), bottom[i]->height(),
			bottom[i]->channels(), coords_, class_num_, 0, 0, 5);
	softmax_cpu(bottom_data + index, class_num_, bottom[i]->num() * anchor_num_,
			bottom[i]->count() / anchor_num_, bottom[i]->count(2), (Dtype)1.0,
			bottom[i]->count(2), (Dtype)1.0, top_data + index);

	// convert coord to image coord
	for (int r = 0; r < bottom[i]->height(); ++r) {
		for (int c = 0; c < bottom[i]->width(); ++c) {
			for (int a = 0; a < anchor_num_; ++a) {
				int i_p = a * (coords_ + 1 + class_num_);

				int index = i_p * bottom[i]->count(2) + r * bottom[i]->width() + c;
				int box_index = index;
				Dtype x = (top_data[box_index] + c) / bottom[i]->width();
				top_data[box_index] = x;
				box_index += bottom[i]->count(2);
				Dtype y = (top_data[box_index] + r) / bottom[i]->height();
				top_data[box_index] = y;
				box_index += bottom[i]->count(2);
				Dtype w = exp(top_data[box_index]) * anchors_[a][0] / bottom[i]->width();
				top_data[box_index] = w;
				box_index += bottom[i]->count(2);
				Dtype h = exp(top_data[box_index]) * anchors_[a][1] / bottom[i]->height();
				top_data[box_index] = h;

				int obj_index = (i_p + coords_) * bottom[i]->count(2) + r * bottom[i]->width() + c;
				Dtype scale = top_data[obj_index];
				top_data[obj_index] = scale;

				for (int l = 0; l < class_num_; ++l) {
					int cla_index = (i_p + coords_ + 1 + l) * bottom[i]->count(2) + r * bottom[i]->width() + c;
					top_data[cla_index] = top_data[cla_index] * scale;
				}
			}
		}
	}
  } // end i
}

INSTANTIATE_CLASS(YOLOv2DetectorLayer);
REGISTER_LAYER_CLASS(YOLOv2Detector);

} // end caffe
